package com.weave.tree;

import java.util.Comparator;

/**
 * @Description
 * 红黑树
 * @Author weave
 * @Date 2023/11/14 20:57
 */
public class RBTree<E> extends BBST<E> {
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    public RBTree(){
        this(null);
    }

    public RBTree(Comparator<E> comparator){
        super(comparator);
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new RBNode<>(element, parent);
    }

    /**
     * 结点默认是红色
     * @param node 新添加的节点
     */
    @Override
    protected void afterAdd(Node<E> node) {
        Node<E> parent = node.parent;

        // 添加的是根节点 或者 上溢到达了根节点
        if(parent == null){
            // 根节点必须是黑的
            black(node);
            return;
        }

        // 如果父节点是黑色的，直接返回
        if(isBlack(parent)){
            return;
        }

        // 叔父结点
        Node<E> uncle = parent.sibling();
        // 祖父结点；上溢需要将这个结点当作一个新添加的红色结点；非上溢需要将这个结点染红
        Node<E> grand = red(parent.parent);
        // uncle是红色，则发生上溢
        if(isRed(uncle)){
            black(parent);
            black(uncle);
            // 把祖父节点当做是新添加的结点
            afterAdd(grand);
            return;
        }

        // 叔父结点不是红色
        if(parent.isLeftChild()){ //L
            if(node.isLeftChild()){ //LL
                black(parent);
            }else { // LR
                black(node);
                rotateLeft(parent);
            }
            rotateRight(grand);
        }else{ //R
            if(node.isLeftChild()){ //RL
                black(node);
                rotateRight(parent);
            }else { //RR
                black(parent);
            }
            rotateLeft(grand);
        }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        // 如果删除的节点是红色
        // 或者 用以取代删除节点的子节点是红色
        if(isRed(node)){
            black(node);
            return;
        }

        Node<E> parent = node.parent;
        // 根节点
        if(parent == null){
            return;
        }

        // 删除的是黑色叶子节点【下溢】
        // 判断被删除的node是左还是右
        boolean left = parent.left == null || node.isLeftChild();
        Node<E> sibling = left ? parent.right : parent.left;
        if (left) { // 被删除的节点在左边，兄弟节点在右边
            if (isRed(sibling)) { // 兄弟节点是红色
                black(sibling);
                red(parent);
                rotateLeft(parent);
                // 更换兄弟
                sibling = parent.right;
            }

            // 兄弟节点必然是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                // 兄弟节点没有1个红色子节点，父节点要向下跟兄弟节点合并
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if (parentBlack) {
                    afterRemove(parent);
                }
            } else { // 兄弟节点至少有1个红色子节点，向兄弟节点借元素
                // 兄弟节点的右边是黑色，兄弟要先旋转
                if (isBlack(sibling.right)) {
                    rotateRight(sibling);
                    sibling = parent.right;
                }

                color(sibling, colorOf(parent));
                black(sibling.right);
                black(parent);
                rotateLeft(parent);
            }
        } else { // 被删除的节点在右边，兄弟节点在左边
            if (isRed(sibling)) { // 兄弟节点是红色
                black(sibling);
                red(parent);
                rotateRight(parent);
                // 更换兄弟
                sibling = parent.left;
            }

            // 兄弟节点必然是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                // 兄弟节点没有1个红色子节点，父节点要向下跟兄弟节点合并
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if (parentBlack) {
                    afterRemove(parent);
                }
            } else { // 兄弟节点至少有1个红色子节点，向兄弟节点借元素
                // 兄弟节点的左边是黑色，兄弟要先旋转
                if (isBlack(sibling.left)) {
                    rotateLeft(sibling);
                    sibling = parent.left;
                }

                color(sibling, colorOf(parent));
                black(sibling.left);
                black(parent);
                rotateRight(parent);
            }
        }
    }

    /**
     * 这种写法，导致AVL树也需要修改接口，所以采用合二为一的方式
     * @param node 要删除的节点
     * @param replacement 取代删除的结点
     */
    protected void afterRemove2(Node<E> node,Node<E> replacement) {
       // 删除的结点是红色
        if(isRed(node)){
            return;
        }

        // 取代结点是红色
        if(isRed(replacement)){
            black(replacement);
            return;
        }

        Node<E> parent = node.parent;
        // 说明node是根节点
        if(parent == null){
            return;
        }

        /*
            到这一步，可以肯定的是，删除的node结点是黑色叶子结点
            注意remove方法，叶子结点的时候，parent的左或者右已经给删除了，所以node.isLeftChild来判断是左还是右是不可以的
            但是因为parent的左或者右是通过叶子结点是左还是右判断的，所以可以如下写法
         */
        boolean left = (parent.left == null) || node.isLeftChild();
        Node<E> sibling = left ? parent.right : parent.left;
        if(left){ //被删除的结点在左边，兄弟结点在右边
            // 这一步的就是让兄弟的儿子(也就是侄子)成为node的兄弟结点
            if(isRed(sibling)){ // 兄弟结点是红色
                black(sibling);
                red(parent);
                rotateLeft(parent);
                // 更换兄弟
                sibling = parent.right;
            }

            // 兄弟节点必然是黑色
            if(isBlack(sibling.left) && isBlack(sibling.right)){
                // 兄弟结点没有1个红色子结点，父节点向下与兄弟结点合并
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if(parentBlack){
                    afterRemove2(parent,null);
                }
            }else{ //兄弟结点至少有1个红色子结点，向兄弟结点借元素
                // 兄弟结点的右边是黑色，兄弟先旋转
                if(isBlack(sibling.right)){
                    rotateRight(sibling);
                    sibling = parent.right;
                }

                color(sibling,colorOf(parent));
                black(sibling.right);
                black(parent);
                rotateLeft(parent);
            }
        }else{ // 被删除的结点在右边，兄弟节点在左边
            if(isRed(sibling)){ // 兄弟结点是红色
                black(sibling);
                red(parent);
                // 这一步，会更新parent的left结点，所以下一步的parent.left不是原来的parent.left。这里感觉是RL
                rotateRight(parent);
                // 更换兄弟
                sibling = parent.left;
            }

            // 兄弟节点必然是黑色
            if(isBlack(sibling.left) && isBlack(sibling.right)){
                // 兄弟结点没有1个红色子结点，父节点向下与兄弟结点合并
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if(parentBlack){
                    afterRemove2(parent,null);
                }
            }else{ //兄弟结点至少有1个红色子结点，向兄弟结点借元素
                // 兄弟结点的左边是黑色，兄弟先旋转
                if(isBlack(sibling.left)){
                    rotateLeft(sibling);
                    sibling = parent.left;
                }

                color(sibling,colorOf(parent));
                black(sibling.left);
                black(parent);
                rotateRight(parent);
            }
        }
    }

    /**
     * 染色
     * @param node
     * @param color
     * @return
     */
    private Node<E> color(Node<E> node,boolean color){
        if (node == null) {
            return null;
        }
        ((RBNode<E>)node).color = color;
        return node;
    }

    private Node<E> red(Node<E> node){
        return color(node,RED);
    }

    private Node<E> black(Node<E> node){
        return color(node,BLACK);
    }

    /**
     * 判断颜色
     * @param node
     * @return
     */
    private boolean colorOf(Node<E> node){
        return node == null ? BLACK : ((RBNode<E>)node).color;
    }

    private boolean isBlack(Node<E> node){
        return colorOf(node) == BLACK;
    }

    private boolean isRed(Node<E> node){
        return colorOf(node) == RED;
    }

    private static class RBNode<E> extends Node<E> {
        boolean color = RED;

        public RBNode(E element, Node<E> parent) {
            super(element, parent);
        }

        @Override
        public String toString() {
            String str = "B_";
            if(color == RED){
                str = "R_";
            }
            return str + element.toString();
        }
    }
}
